# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openFuyao is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
This module defines NPU collector error detection rules, including HBM, DDR,
NPU core, and network error rules. Each rule collects collector metrics via
MainCollector, evaluates them against thresholds or error codes, and generates
HardwareEvent instances for further fault handling and recovery.
"""

import logging
import os
import random
from typing import List, Optional

import hardware_diagnosis.common.constant as const
from hardware_diagnosis.detector.base import ErrorRule
from hardware_diagnosis.event.hardware_event import NPUHardwareEvent
from hardware_diagnosis.event.base import FaultType
from hardware_diagnosis.recovery.hardware_strategies.npu.utils import get_device_error_code_string

module_name = os.getenv("MODULE_NAME", "default")

logger = logging.getLogger(module_name)


class HBMErrorRule(ErrorRule):
    """
    HBMErrorRule detects HBM-related errors such as over-temperature and reported error events.
    """
    def __init__(self):
        super().__init__("NPU_HBM")
        self.hbm_resource = []

    def error_code(self, event: dict) -> str:
        return str(event["error_code"])

    def detect(self, data: dict) -> Optional[List[NPUHardwareEvent]]:
        events = []
        hbm_data = data.get("collector.npu.hbm", {})
        for logic_id, metrics in hbm_data.items():
            node_name = metrics.get("node_name")
            card_id = metrics.get("card_id")
            device_id = metrics.get("device_id")
            self.hbm_resource.append(metrics)
            temperature = metrics.get("npu_chip_info_hbm_temperature")
            if (temperature != "NA" and int(temperature)) > const.HBM_TEMPERATURE_THRESHOLD:
                event = NPUHardwareEvent(hardware_type=self.name,
                                         error_code=hex(100),
                                         node_name=node_name,
                                         card_id=card_id,
                                         device_id=device_id,
                                         logic_id=logic_id,
                                         fault_type=FaultType.HARDWARE,
                                         details="HBM temperature is over threshold")
                events.append(event)

            if metrics.get("errors"):
                for error_event in metrics["errors"]:
                    event = NPUHardwareEvent(hardware_type=self.name,
                                             error_code=self.error_code(error_event),
                                             node_name=node_name,
                                             card_id=card_id,
                                             device_id=device_id,
                                             logic_id=logic_id,
                                             fault_type=FaultType.HARDWARE,
                                             details=error_event["error_message"])
                    events.append(event)
        return events


class DDRErrorRule(ErrorRule):
    """
    DDRErrorRule detects DDR-related error events from collected metrics.
    """
    def __init__(self):
        super().__init__("NPU_DDR")
        self.ddr_resource = []

    def error_code(self, event) -> str:
        return str(event["error_code"])

    def detect(self, data: dict) -> Optional[List[NPUHardwareEvent]]:
        events = []
        ddr_data = data.get("collector.npu.ddr", {})
        for logic_id, metrics in ddr_data.items():
            node_name = metrics.get("node_name")
            device_id = metrics.get("device_id")
            card_id = metrics.get("card_id")
            self.ddr_resource.append(metrics)
            if metrics.get("errors"):
                for error_event in metrics["errors"]:
                    event = NPUHardwareEvent(hardware_type=self.name,
                                             error_code=self.error_code(error_event),
                                             node_name=node_name,
                                             card_id=card_id,
                                             device_id=device_id,
                                             logic_id=logic_id,
                                             fault_type=FaultType.HARDWARE,
                                             details=error_event["error_message"])
                    events.append(event)
        return events


class NPUErrorRules(ErrorRule):
    """
    NPUErrorRules detects generic NPU core errors based on error codes and reported faults.
    """
    def __init__(self):
        super().__init__("NPU")
        self.npu_resource = []

    def error_code(self, event) -> str:
        return str(event["error_code"])

    def detect(self, data: dict) -> Optional[List[NPUHardwareEvent]]:
        events = []
        npu_data = data.get("collector.npu.npu", {})

        # test
        # error_code_samples = random.sample(list(const.ERROR_CODE_MAP.keys()), 2)
        # error_code_samples_int = [int(code, 16) for code in error_code_samples]
        # test_data = {
        #     "node_name": "test",
        #     "device_id": 101,
        #     "card_id": 101,
        #     "npu_chip_info_error_code": error_code_samples_int
        #     }
        # npu_data[101] = test_data

        for logic_id, metrics in npu_data.items():
            node_name = metrics.get("node_name")
            device_id = metrics.get("device_id")
            card_id = metrics.get("card_id")
            self.npu_resource.append(metrics)
            if metrics.get("npu_chip_info_error_code"):
                for errcode in metrics["npu_chip_info_error_code"]:
                    try:
                        detail, _, err = get_device_error_code_string(logic_id, card_id, device_id,
                                                                      errcode)
                        event = NPUHardwareEvent(hardware_type=self.name,
                                                 error_code=hex(errcode),
                                                 node_name=node_name,
                                                 card_id=card_id,
                                                 device_id=device_id,
                                                 logic_id=logic_id,
                                                 fault_type=FaultType.HARDWARE,
                                                 details=detail)
                        events.append(event)
                    except Exception as e:  # pylint: disable=broad-except
                        logger.error("Failed to obtain NPU error message: %s", {str(e)})
            if metrics.get("errors"):
                for err in metrics["errors"]:
                    event = NPUHardwareEvent(hardware_type=self.name,
                                             error_code=self.error_code(err),
                                             node_name=node_name,
                                             card_id=card_id,
                                             device_id=device_id,
                                             logic_id=logic_id,
                                             fault_type=FaultType.HARDWARE, details=err["error_message"])
                    events.append(event)
        return events


class NetworkErrorRule(ErrorRule):
    """
    NetworkErrorRule detects network-related errors reported by NPU metrics.
    """
    def __init__(self):
        super().__init__("NPU_Network")
        self.network_resource = []

    def error_code(self, event) -> str:
        return str(event["error_code"])

    def detect(self, data: dict) -> Optional[List[NPUHardwareEvent]]:
        events = []
        network_data = data.get("collector.npu.network", {})
        for logic_id, metrics in network_data.items():
            node_name = metrics.get("node_name")
            device_id = metrics.get("device_id")
            card_id = metrics.get("card_id")
            self.network_resource.append(metrics)
            if metrics.get("errors"):
                for error_event in metrics["errors"]:
                    event = NPUHardwareEvent(hardware_type=self.name,
                                             error_code=self.error_code(error_event),
                                             node_name=node_name,
                                             card_id=card_id,
                                             device_id=device_id,
                                             logic_id=logic_id,
                                             fault_type=FaultType.HARDWARE,
                                             details=error_event["error_message"])
                    events.append(event)
        return events
